package org.jeecg.common.util.sm.utils;

import org.bouncycastle.asn1.*;
import org.bouncycastle.asn1.cms.*;
import org.bouncycastle.asn1.x509.*;
import org.bouncycastle.cms.CMSEnvelopedData;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.digests.SM3Digest;
import org.bouncycastle.crypto.engines.SM2Engine;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.crypto.params.ECKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.crypto.params.ParametersWithRandom;
import org.bouncycastle.jcajce.provider.asymmetric.util.ECUtil;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.util.encoders.Base64;
import org.jeecg.common.util.sm.bean.KeyAndIV;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.security.InvalidKeyException;
import java.security.Security;
import java.security.cert.CertificateEncodingException;
import java.security.cert.X509Certificate;


public class SM2Utils {

    static {
        Security.addProvider(new BouncyCastleProvider());
    }

    public static String sm4Aly = "SM4/CBC/PKCS7Padding";
    public static String SM2_pubKey_encrypt_oid = "1.2.156.10197.1.301.3";
    public static String sm2Data_oid = "1.2.156.10197.6.1.4.2.1";
    public static String sm2EnvelopedData_oid = "1.2.156.10197.6.1.4.2.3";
    public static String sm4_CBC_oid = "1.2.156.10197.1.104.2";
    public static String SubjectKey_oid = "2.5.29.14";


    public static String envelopData(X509Certificate cert, String source) {
        try {
            //生成对称密钥key和IV
            KeyAndIV keyAndIV = KeyAndIV.generateKeyAndIV(16, 16);
            //创建SM2Engine
            SM2Engine engine = new SM2Engine(new SM3Digest(), SM2Engine.Mode.C1C3C2);
            ECPublicKeyParameters ecPublicKeyParameters = (ECPublicKeyParameters) ECUtil.generatePublicKeyParameter(cert.getPublicKey());
            ParametersWithRandom parametersWithRandom = new ParametersWithRandom(ecPublicKeyParameters);
            engine.init(true, parametersWithRandom);
            //使用公钥对对称密钥加密
            byte[] encryptedBytes = engine.processBlock(keyAndIV.getKey(), 0, keyAndIV.getKey().length);

            //构造sm2加密结果结构体
            ASN1OctetString sm2CipherObject = getSm2CipherObject(encryptedBytes, parametersWithRandom);

            //使用对称密钥对原文进行加密
            byte[] sm4EncryptCBCByte = SM4KeyUtils.getSm4Encrypt(sm4Aly, source, keyAndIV);
            ASN1OctetString encryptedOctet = new DEROctetString(sm4EncryptCBCByte);

            EncryptedContentInfo encryptedContentInfo = new EncryptedContentInfo(new ASN1ObjectIdentifier(sm2Data_oid),
                    new AlgorithmIdentifier(new ASN1ObjectIdentifier(sm4_CBC_oid), new DEROctetString(keyAndIV.getIv())), encryptedOctet);
            //获取subjectKeyIdentifier 内容
            RecipientIdentifier recipientIdentifier = getRecipientIdentifier(cert);
            //构造接收者信封结构体
            ASN1EncodableVector recipientInfos = getAsn1EncodableVector(recipientIdentifier, sm2CipherObject);

            EnvelopedData envelopedData = new EnvelopedData(null, new DERSet(recipientInfos), encryptedContentInfo, ASN1Set.getInstance(null));
            ContentInfo contentInfo = new ContentInfo(new ASN1ObjectIdentifier(sm2EnvelopedData_oid), envelopedData);
            CMSEnvelopedData cmsEnvelopedData = new CMSEnvelopedData(contentInfo);
            return Base64.toBase64String(cmsEnvelopedData.getEncoded());
        } catch (InvalidKeyException | InvalidCipherTextException | CertificateEncodingException | IOException |
                 CMSException e) {
            e.printStackTrace();
            return "";
        }
    }

    private static ASN1EncodableVector getAsn1EncodableVector(RecipientIdentifier recipientIdentifier, ASN1OctetString sm2CipherObject) {
        KeyTransRecipientInfo keyTransRecipientInfo = new KeyTransRecipientInfo(recipientIdentifier, new AlgorithmIdentifier(new ASN1ObjectIdentifier(SM2_pubKey_encrypt_oid), DERNull.INSTANCE), sm2CipherObject);
        RecipientInfo receiverCerts = new RecipientInfo(keyTransRecipientInfo);
        ASN1EncodableVector recipientInfos = new ASN1EncodableVector();
        recipientInfos.add(receiverCerts);
        return recipientInfos;
    }

    /**
     * 从证书中解析subjectKeyIdentifier
     */
    private static RecipientIdentifier getRecipientIdentifier(X509Certificate cert) throws CertificateEncodingException, IOException {
        TBSCertificate instance = TBSCertificate.getInstance(cert.getTBSCertificate());
        Extensions extensions = instance.getExtensions();
        Extension extension = extensions.getExtension(new ASN1ObjectIdentifier(SubjectKey_oid));
        byte[] octets = extension.getExtnValue().getOctets();
        ASN1InputStream asn1InputStream = new ASN1InputStream(new ByteArrayInputStream(octets));
        ASN1Object asn1Object = asn1InputStream.readObject();
        SubjectKeyIdentifier subjectKeyIdentifier = SubjectKeyIdentifier.getInstance(asn1Object);
        DEROctetString derOctetString = new DEROctetString(subjectKeyIdentifier.getKeyIdentifier());
        return new RecipientIdentifier(derOctetString);
    }

    /**
     * 根据加密结果构建需要解密内容结构体
     */

    private static ASN1OctetString getSm2CipherObject(byte[] bytes, ParametersWithRandom parametersWithRandom) throws IOException {

        ECKeyParameters ecKey = (ECKeyParameters) parametersWithRandom.getParameters();
        ECDomainParameters ecParams = ecKey.getParameters();

        byte[] c1Bytes = new byte[65];
        System.arraycopy(bytes, 0, c1Bytes, 0, c1Bytes.length);
        byte[] c3Bytes = new byte[32];
        System.arraycopy(bytes, c1Bytes.length, c3Bytes, 0, c3Bytes.length);
        byte[] c2Bytes = new byte[bytes.length - 32 - 65];
        System.arraycopy(bytes, c1Bytes.length + c3Bytes.length, c2Bytes, 0, c2Bytes.length);
        ECPoint c1Pont = ecParams.getCurve().decodePoint(c1Bytes);
        BigInteger X = c1Pont.normalize().getAffineXCoord().toBigInteger();
        BigInteger Y = c1Pont.normalize().getAffineYCoord().toBigInteger();

        ASN1EncodableVector v = new ASN1EncodableVector();
        v.add(new ASN1Integer(X));
        v.add(new ASN1Integer(Y));
        v.add(new DEROctetString(c3Bytes));
        v.add(new DEROctetString(c2Bytes));
        return new DEROctetString(new DERSequence(v));
    }
}
